from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch
from flask import Flask, request

# Load OpenAI GPT-2 model from local
model_path = "C:/01.project/chatgpt/gpt-2/models/124M"
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained(model_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

app = Flask(__name__)

# Define API endpoint
@app.route('/generate_text', methods=['GET'])
def generate_text():
    prompt = request.args.get('prompt', '')
    length = int(request.args.get('length', '50'))
    input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
    output = model.generate(input_ids=input_ids, max_length=length, do_sample=True, temperature=1.0)
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return json.dumps({'text': generated_text})

if __name__ == '__main__':
    app.run()